from .op_basic import *
from ..utils import get_default_manager
import pandas as pd


def select_columns(instance, column_list, **kwargs):
    default_manager = get_default_manager()
    overwrite_default_helper(
        get_default_name_kwargs(default_manager, 'select_columns'),
        kwargs,
    )

    def __select_columns(manager, sess, df):
        return df[column_list]

    return Instance(__select_columns, [instance], manager=default_manager, **kwargs)


@as_op('select_conc_columns')
def select_conc_columns(df, must_select=[], must_drop=[], conc_max_value_num=15):
    '''
    根据这一列特征的个数判断, 选择出可能是离散x的列。
    must_select里面的列必然被选上，一般是把y放在里面！
    must_drop里面的列必然不被选上
    conc_max_value_num是作为离散变量，最多有那么多个值。多于conc_max_value_num的值被认为是连续变量。

    返回选择后的dataframe
    '''
    assert (df.columns.duplicated() == False).all()
    selected = []
    for col in df.columns:
        if ((len(df[col].unique()) <= conc_max_value_num) or (col in must_select)) and (col not in must_drop):
            selected.append(col)
    return df[selected]


@as_op('select_cont_columns')
def select_cont_columns(df, must_select=[], must_drop=[], conc_max_value_num=15):
    '''
    根据这一列特征的个数判断, 选择出可能是连续x的列。
    must_select里面的列必然被选上，一般是把y放在里面！
    conc_max_value_num是作为离散变量，最多有那么多个值。多于conc_max_value_num的值被认为是连续变量（和select_conc_columns统一）。

    返回选择后的dataframe
    '''
    assert (df.columns.duplicated() == False).all()
    selected = []
    for col in df.columns:
        if ((len(df[col].unique()) > conc_max_value_num) or (col in must_select)) and (col not in must_drop):
            selected.append(col)
    return df[selected]


@as_op('drop')
def drop(df, *args, **kwargs):
    return df.drop(*args, **kwargs)


@as_op('drop')
def filter(df, *args, **kwargs):
    return df.filter(*args, **kwargs)


@as_op('add_prefix_columns')
def add_prefix_columns(df, prefix, inplace=False):
    assert (df.columns.duplicated() == False).all()
    __replace_dict = {}
    for key in df.columns:
        __replace_dict[key] = prefix + '_' + key
    if inplace:
        df.rename(columns=__replace_dict, inplace=True)
    else:
        return df.rename(columns=__replace_dict, inplace=False)
